namespace NBody.InitialConditions
{
    using System;
    using Fftw;
    using Cosmology;
    
    public struct Displacements
    {
        public double Sx;
        public double Sy;
        public double Sz;
        
        public double SVx;
        public double SVy;
        public double SVz;
    }   
    
    public class Perturbations
    {
        int min_wave;
        int max_wave;
        RandomGaussian random;
        Spectrum cosmology;
        
        public Perturbations(Spectrum cosmo, int seed)
        {
            if (seed != 0)
                random = new RandomGaussian(seed);
            else
                random = new RandomGaussian();
            this.cosmology = cosmo;
            
            min_wave = 0;
            max_wave = -1;
        }
        
        public int MinWave
        {
            get { return min_wave; }
            set { min_wave = value; }
        }
        
        public int MaxWave
        {
            get { return max_wave; }
            set { max_wave = value; }
        }
        
        public Displacements[] CalcGaussDispFFT(int n_per_side, double SideLength)
        {
            Console.Error.WriteLine("Calculating Gaussian displacements using FFT method");
            
            if (max_wave > n_per_side || max_wave < 0)
                max_wave = n_per_side;
            if (min_wave > max_wave || min_wave > n_per_side)
                min_wave = 0;
            
            double delta_k = 2.0 * Math.PI / SideLength;
            
            Complex[] fx = new Complex[n_per_side * n_per_side * (n_per_side/2 + 1)];
            Complex[] fy = new Complex[n_per_side * n_per_side * (n_per_side/2 + 1)];
            Complex[] fz = new Complex[n_per_side * n_per_side * (n_per_side/2 + 1)];
            
            double a, b;
            double kx, ky, kz;
            int c = 0;
            for (int i = 0; i < n_per_side; i++)
            {
                kx = delta_k * i;
                for (int j = 0; j < n_per_side; j++)
                {
                    ky = delta_k * j;
                    for (int k = 0; k < n_per_side/2 + 1; k++)
                    {
                        kz = delta_k * k;
                        double kk = Math.Sqrt(kx*kx + ky*ky + kz*kz);
                        
                        double r1 = random.NextDouble();
                        double r2 = random.NextDouble();
                        
                        if (kk == 0.0)
                        {
                            a = 0.0;
                            b = 0.0;
                        }
                        else if (i < min_wave || j < min_wave || k < min_wave ||
                                 i > max_wave || j > max_wave || k > max_wave)
                        {
                            a = 0.0;
                            b = 0.0;
                        }
                        else
                        {
                            double p = cosmology.PowerSpectrum(kk);// * delta_k * delta_k * delta_k;
                            a = Math.Sqrt(p) / kk / kk * r1;
                            b = Math.Sqrt(p) / kk / kk * r2;
                        }
                        
                        fx[c].Re = 0.5 * kx * b;
                        fx[c].Im = 0.5 * kx * a;
                        
                        fy[c].Re = 0.5 * ky * b;
                        fy[c].Im = 0.5 * ky * a;
                        
                        fz[c].Re = 0.5 * kz * b;
                        fz[c].Im = 0.5 * kz * a;
                        
                        c++;
                    }
                }
            }    
            
            PlanComplexToReal3D planx = new PlanComplexToReal3D(n_per_side, n_per_side, n_per_side, fx, Plan.FFTW_ESTIMATE);
            planx.Execute();
            double[] sx = planx.RealOutput;
            
            PlanComplexToReal3D plany = new PlanComplexToReal3D(n_per_side, n_per_side, n_per_side, fy, Plan.FFTW_ESTIMATE);
            plany.Execute();
            double[] sy = plany.RealOutput;
            
            PlanComplexToReal3D planz = new PlanComplexToReal3D(n_per_side, n_per_side, n_per_side, fz, Plan.FFTW_ESTIMATE);
            planz.Execute();
            double[] sz = planz.RealOutput;
            
            int ntotal = n_per_side * n_per_side * n_per_side;
            Displacements[] disp = new Displacements[ntotal];
            
            double max = 0.0;
            double avg = 0.0;
            double max_s = 0.0;
            double avg_s = 0.0;
            for (int i = 0; i < ntotal; i++)
            {
                double d = Math.Sqrt(sx[i]*sx[i] + sy[i]*sy[i] + sz[i]*sz[i]);
                if (d > max) max = d;
                avg += d;
            }
            avg /= (double)ntotal;
            double grid = SideLength / (double)n_per_side;
                        
            Console.Error.WriteLine("  Grid spacing: {0}", grid);
            Console.Error.WriteLine("  Average displacements: {0}", avg);
            Console.Error.WriteLine("  Max displacement: {0}", max);
            double scale = grid / max;
            double z = 1.0 / Math.Sqrt(scale) - 1.0;
            Console.Error.WriteLine("  Scaling to a redshift of z = {0}", z);
            
            for (int i = 0; i < ntotal; i++)
            {
                disp[i].Sx = cosmology.Model.LinearGrowth(z) * cosmology.Model.LinearGrowth(z) * sx[i];
                disp[i].Sy = cosmology.Model.LinearGrowth(z) * cosmology.Model.LinearGrowth(z) * sy[i];
                disp[i].Sz = cosmology.Model.LinearGrowth(z) * cosmology.Model.LinearGrowth(z) * sz[i];
                
                disp[i].SVx = cosmology.Model.LinearGrowthDerivative(z) * cosmology.Model.LinearGrowth(z) *  sx[i];
                disp[i].SVy = cosmology.Model.LinearGrowthDerivative(z) * cosmology.Model.LinearGrowth(z) *  sy[i];
                disp[i].SVz = cosmology.Model.LinearGrowthDerivative(z) * cosmology.Model.LinearGrowth(z) *  sz[i];
                
                double d_s = Math.Sqrt(disp[i].Sx*disp[i].Sx + disp[i].Sy*disp[i].Sy + disp[i].Sz*disp[i].Sz);
                if (d_s > max_s) max_s = d_s;
                avg_s += d_s;
            }
            avg_s /= (double)ntotal;
            
            Console.Error.WriteLine("    Average scaled displacements: {0}", avg_s);
            Console.Error.WriteLine("    Max scaled displacement: {0}", max_s);
            
            return disp;
        }
    }
}
        
